from . import variable as va
from PIL import Image
import numpy as np
from . import hogfeature as hog
import joblib
import os

'''
===============
    支持函数
===============
'''
# 对传入的单张图片提取特征
def get_one_feat(image_path):
    # 加载图片
    img = Image.open(os.path.join(image_path))
    # 处理图片
    img = hog.dealImage(img)
    # 提取特征
    feat = hog.get_feat(img)
    feat = np.concatenate((feat, []))
    return feat

# 用模型进行分类
def do_Classify(image_path, modelName):
    try:
        # 加载模型
        clf = joblib.load(va.model_path + modelName)
    except:
        print('加载模型失败！请检查……')
    else:
        # 提取图片特征
        feat = get_one_feat(image_path)
        # 进行预测
        result = clf.predict([feat])
        return result

'''
===============
    API函数
===============
'''
# 用SVM模型进行分类
def do_SVM_Classify(image_path):
    result = do_Classify(image_path, 'SVM_model')
    return result
    
# 用KNN模型进行分类
def do_KNN_Classify(image_path):
    result = do_Classify(image_path, 'KNN_model')
    return result

# 用KNN模型进行分类
def do_RF_Classify(image_path):
    result = do_Classify(image_path, 'RF_model')
    return result


if __name__ == '__main__':
    # 加载图片
    image_path = input('请输入图片的路径:')
    # 进行分类
    result = do_SVM_Classify(image_path)
    print(va.labels[int(result)])
